# coding : UTF-8
"""
作者：BingBO   时间：2022年07月24日
自动调整代码格式 ：Alt+Ctrl+L
"""

import cv2 as cv
import numpy as np


# 从棋盘格图像中返回分割后的光条阈值图像（单通道）
def laser_threshold_image_from_checkerboard(image):
    # height, width, channels = img.shape
    _, greenChannel, _ = cv.split(image)
    # cv.imshow('greenChannel', greenChannel)
    # maxVaL = np.max(greenChannel)
    # 找出合适的阈值，这里默认是最大值255，也能够取得比较好的结果和普遍性
    maxVaL = 255
    _, thresholdImg = cv.threshold(greenChannel, maxVaL - 10, 255, cv.THRESH_BINARY)
    thresholdImg = cv.medianBlur(thresholdImg, 3)
    # cv.imshow('threshold', thresholdImg)
    return thresholdImg


# (垂直)灰度重心法
def center_of_gravity(laser_threshold_image):
    hight, width = laser_threshold_image.shape
    center_point_coordinates = []
    for col_idx in range(width):
        n = 0
        row_idx_sum = 0
        for row_idx in range(hight):
            if laser_threshold_image[row_idx, col_idx] != 0:
                n += 1
                row_idx_sum += row_idx + 1
        if n:
            gravity_row_idx = round(row_idx_sum / n) - 1
            center_point_coordinates.append((col_idx, gravity_row_idx))
    return center_point_coordinates


# 最小二乘法拟合这些中心坐标直线，并查看拟合显示效果
def laser_polyfit(center_point_coordinates, img):
    x_list = []
    y_list = []
    for item in center_point_coordinates:
        x_list.append(item[0])
        y_list.append(item[1])
    k, b = np.polyfit(x_list, y_list, 1)  # 用最小二乘法（含COG_V误差）拟合得到的直线参数
    print("一阶拟合系数为:", k, b)
    # A, B, C = k, -1, b

    # 在图像上显示拟合线条
    if k != 0:
        x_left_point = center_point_coordinates[0][0]
        x_right_point = center_point_coordinates[-1][0]
        y_left_point = k * x_left_point + b
        y_right_point = k * x_right_point + b
        # 绘制焊缝线
        cv.line(img, (x_left_point, round(y_left_point)), (x_right_point, round(y_right_point)), (255, 255, 0), 1)  # 青蓝
        # cv.imshow('result', img)
    return k, b  # [A, B, C]


# 用户接口
def get_laser_polyfit_from_checkerboard(src_img):
    laser_threshold_image = laser_threshold_image_from_checkerboard(src_img)
    center_point_coordinates = center_of_gravity(laser_threshold_image)
    k, b = laser_polyfit(center_point_coordinates, src_img)
    return k, b


if __name__ == "__main__":
    src = cv.imread(r"E:\SHU\Research Group\LaserVisionSensor\Calibrate\Image_20231110134340816.jpg")
    src = cv.resize(src, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)

    t1 = cv.getTickCount()
    get_laser_polyfit_from_checkerboard(src)
    t2 = cv.getTickCount()
    print("time: %s ms" % ((t2 - t1) / cv.getTickFrequency() * 1000))

    cv.waitKey(0)
    cv.destroyAllWindows()
